"""
建立ARMA和ARIMA模型
"""
import pandas as pd
import numpy as np
from statsmodels.tsa.seasonal import seasonal_decompose
import statsmodels.tsa.stattools as ts
from statsmodels.tsa.arima_model import ARIMA
from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt

plt.style.use(['bmh'])


def dateparse(x): return pd.datetime.strptime(x, '%Y-%m')


data = pd.read_csv('data/AirPassengers.csv',
                   parse_dates=['Month'], date_parser=dateparse)
data = data.set_index('Month')

# 将数据分解为趋势序列，季节序列和残差序列
ts_log = np.log(data['Passengers'])
decomposition = seasonal_decompose(ts_log, freq=12)
trend = decomposition.trend  # 趋势
seasonal = decomposition.seasonal  # 季节性
residual = decomposition.resid  # 残差序列
residual.dropna(inplace=True)
# 判断残差序列的稳定性：
dftest = ts.adfuller(residual)
dfoutput = pd.Series(dftest[0:4], index=[
                     'Test Statistic', 'p-value', '#Lags Used', 'Number of Obserfvisions Used'])
for key, value in dftest[4].items():
    dfoutput['Critical Value (%s)' % key] = value
print(dfoutput)
# 用ARIMA方法来进行预测
model_ARIMA = ARIMA(residual, (2, 0, 2)).fit(disp=-1, method='css')
predictions_ARIMA = model_ARIMA.predict(start='1950-01', end='1962-04')
# 转换回原数据空间
predictions_ARIMA = predictions_ARIMA.add(
    trend, fill_value=0).add(seasonal, fill_value=0)
predictions_ARIMA = np.exp(predictions_ARIMA)
# 用线性拟合来预测1960-06之后的数据
trend.dropna(inplace=True)
x = pd.Series(range(trend.size), index=trend.index).to_frame()
regsr = LinearRegression().fit(x, trend)
x = pd.Series(range(0, 154), index=(
    pd.period_range('1949-07', periods=154, freq='M')))
x = x.to_frame()
res_predict = regsr.predict(x)
trend_fitted = pd.Series(res_predict, index=x.index).to_timestamp()
# 扩展seasonal的值到1962-04
index_shift = pd.period_range('1949-01', periods=160, freq='M').to_timestamp()
seasonal = seasonal.reindex(index_shift)
seasonal_shifted = seasonal.shift(24)
# 用新的trend和seasonal来返回到原数据
model_ARIMA_new = ARIMA(residual, (2, 0, 2)).fit(disp=-1, method='css')
predictions_ARIMA_new = model_ARIMA_new.predict(start='1950-01', end='1962-04')
predictions_ARIMA_new = predictions_ARIMA_new.add(
    trend_fitted, fill_value=0).add(seasonal_shifted, fill_value=0)
predictions_ARIMA_new = np.exp(predictions_ARIMA_new)

# 绘图
plt.figure(figsize=(18, 7), dpi=128)

plt.subplot(121)
plt.title('Time Series for AirPassengers')
plt.xlabel('Month')
plt.ylabel('Passengers')
plt.plot(data['Passengers'], 'b-', label='Origin')
plt.plot(predictions_ARIMA, 'r-', label='ARIMA Prediction')
plt.plot(predictions_ARIMA_new, 'g--', label='ARIMA Prediction_New')
plt.legend(loc='upper right')

plt.subplot(122)
plt.title('Trend & Seasonal')
plt.plot(trend, label='Trend')
plt.plot(trend_fitted, label='Trend_Pred')
plt.plot(seasonal, label='Seasonal')
plt.plot(seasonal_shifted, label='Seasonal_Shift')
plt.legend(loc='center right')

plt.savefig('第8章：时间序列/Time Series')
